import pandas as pd
import numpy as np

def span_df_by_seconds(df, trucks):
    '''
    Input:
    df: start_time, end_time, status, truck_id
    trucks
    Output:
    df_span: ts, status, truck_id
    '''
    df_span = pd.DataFrame()
    df_span = []
    for truck in trucks:
        df_tmp = df.loc[df['truck_id'] == truck]
        for k, row in df_tmp.iterrows():
            start = row['start_time']
            end = row['end_time']
            ts = pd.date_range(start=start, end=end, freq='S')
            sta = row['status']
            spaned = pd.DataFrame({'ts': ts, 'status':sta, 'truck_id':truck})
            spaned = spaned.to_numpy().tolist()
            # df_span = pd.concat([df_span, spaned], axis=0)
            df_span.extend(spaned)

    df_span = pd.DataFrame(df_span, columns=['ts', 'status', 'truck_id'])
    df_span = df_span.sort_values(by=['truck_id', 'ts']).reset_index(drop=True)
    return df_span
def compress_df(df_span, trucks, status_col = None):
    '''
    Input: 
    df_span: ts, status, truck_id
    trucks:
    status_col: 哪一列列名是关注的状态
    Output:
    result_df: start_time, end_time, status, truck_id, duration_time
    '''
    # result_df = pd.DataFrame(columns=['start_time', 'end_time', 'status', 'lat_start', 'lat_end', 
    #                               'lon_start', 'lon_end', 'truck_id'])
    if status_col is not None:
        df_span.rename(columns={status_col:'status'})

    df_span['judge'] = df_span['status'].astype(str) + df_span['module_name'].astype(str) + df_span['item_name'].astype(str) + df_span['level'].astype(str)
    result_df = []
    for truck in trucks:
        tmp_df = df_span.loc[df_span['truck_id'] == truck].reset_index(drop=True)
        speed = []
        latitude = []
        longitude = []
        sta_ready = tmp_df['status'].loc[0]
        module_name_ready = tmp_df['module_name'].loc[0]
        item_name_ready = tmp_df['item_name'].loc[0]
        level_ready = tmp_df['level'].loc[0]
        judge_ready = tmp_df['judge'].loc[0]
        lat_start_ready = tmp_df['latitude'].loc[0]
        lon_start_ready = tmp_df['longitude'].loc[0]
        start_time = tmp_df['ts'].loc[0]
        speed.append(tmp_df['speed'].loc[0])
        latitude.append(tmp_df['latitude'].loc[0])
        longitude.append(tmp_df['longitude'].loc[0])
        for k in range(1, len(tmp_df)):
            if k == tmp_df.index[-1]:
                if tmp_df['judge'].loc[k] == judge_ready:
                    end_time = tmp_df['ts'].loc[k]
                    status = sta_ready
                    module_name = module_name_ready
                    item_name = item_name_ready
                    level = level_ready
                    lan_start = lat_start_ready
                    lon_start = lon_start_ready
                    lat_end =  tmp_df['latitude'].loc[k]
                    lon_end =  tmp_df['longitude'].loc[k]
                    speed.append(tmp_df['speed'].loc[k])
                    latitude.append(tmp_df['latitude'].loc[k])
                    longitude.append(tmp_df['longitude'].loc[k])
                    row_df = [start_time, end_time, status, module_name, item_name, level,
                              lan_start, lon_start, lat_end, lon_end, str(latitude), str(longitude) ,str(speed), truck]
                    # row_df = pd.DataFrame({'start_time': start_time, 'end_time': end_time, 'status': status, 'truck_id':truck}, index=[0])
                    # result_df = pd.concat([result_df, row_df], axis=0)
                    result_df.append(row_df)
                else:
                    end_time = tmp_df['ts'].loc[k-1]
                    status = sta_ready
                    module_name = module_name_ready 
                    item_name = item_name_ready
                    level = level_ready
                    lan_start = lat_start_ready
                    lon_start = lon_start_ready
                    lat_end =  tmp_df['latitude'].loc[k-1]
                    lon_end =  tmp_df['longitude'].loc[k-1]
                    row_df = [start_time, end_time, status, module_name, item_name, level, 
                              lan_start, lon_start, lat_end, lon_end, str(latitude), str(longitude), str(speed),truck]
                    result_df.append(row_df)
                    sta_ready = tmp_df['status'].loc[k]
                    module_name_ready = tmp_df['module_name'].loc[k]
                    item_name_ready = tmp_df['item_name'].loc[k]
                    level_ready = tmp_df['level'].loc[k]
                    judge_ready = tmp_df['judge'].loc[k]
                    lat_start_ready = tmp_df['latitude'].loc[k]
                    lon_start_ready = tmp_df['longitude'].loc[k]
                    speed = []
                    latitude = []
                    longitude = []
                    speed.append(tmp_df['speed'].loc[k])
                    latitude.append(tmp_df['latitude'].loc[k])
                    longitude.append(tmp_df['longitude'].loc[k])

                    start_time = tmp_df['ts'].loc[k]
                    end_time = tmp_df['ts'].loc[k]
                    status = sta_ready
                    module_name = module_name_ready
                    item_name = item_name_ready
                    level = level_ready
                    lan_start = lat_start_ready
                    lon_start = lon_start_ready
                    lat_end =  tmp_df['latitude'].loc[k]
                    lon_end =  tmp_df['longitude'].loc[k]
                    
                    row_df = [start_time, end_time, status, module_name, item_name, level, 
                              lan_start, lon_start, lat_end, lon_end, str(latitude), str(longitude), str(speed), truck]
                    # row_df = pd.DataFrame({'start_time': start_time, 'end_time': end_time, 'status': status, 'truck_id':truck}, index=[0])
                    # result_df = pd.concat([result_df, row_df], axis=0)
                    result_df.append(row_df)
                continue
            if tmp_df['judge'].loc[k] == judge_ready and ((tmp_df['ts'].loc[k] - tmp_df['ts'].loc[k-1]).total_seconds() <= 120):
                speed.append(tmp_df['speed'].loc[k])
                latitude.append(tmp_df['latitude'].loc[k])
                longitude.append(tmp_df['longitude'].loc[k])                
                continue
            else:
                end_time = tmp_df['ts'].loc[k-1]
                status = sta_ready
                module_name = module_name_ready
                item_name = item_name_ready
                level = level_ready
                lan_start = lat_start_ready
                lon_start = lon_start_ready
                lat_end =  tmp_df['latitude'].loc[k-1]
                lon_end =  tmp_df['longitude'].loc[k-1]
                row_df = [start_time, end_time, status, module_name, item_name, level, 
                          lan_start, lon_start, lat_end, lon_end, str(latitude), str(longitude), str(speed), truck]
                result_df.append(row_df)
                sta_ready = tmp_df['status'].loc[k]
                module_name_ready = tmp_df['module_name'].loc[k]
                item_name_ready = tmp_df['item_name'].loc[k]
                level_ready = tmp_df['level'].loc[k]
                judge_ready = tmp_df['judge'].loc[k]
                start_time = tmp_df['ts'].loc[k]
                lat_start_ready = tmp_df['latitude'].loc[k]
                lon_start_ready = tmp_df['longitude'].loc[k]
                speed = []
                latitude = []
                longitude = []
                speed.append(tmp_df['speed'].loc[k])
                latitude.append(tmp_df['latitude'].loc[k])
                longitude.append(tmp_df['longitude'].loc[k])
    result_df = pd.DataFrame(result_df, columns=['start_time', 'end_time', 'status', 'module_name',
                                                 'item_name', 'level', 'lat_start', 'lon_start', 'lat_end', 'lon_end', 'latitude',
                                                    'longitude', 'speed', 'truck_id'])
    result_df['duration_time'] = (result_df['end_time'] - result_df['start_time']).dt.total_seconds()

    return result_df
            
